from environment.trace_sample import Trace
from environment.rewards import Reward
from config import DRLParameters
import numpy as np
import copy


class State:
    def __init__(self, trace: Trace, window_size: int, dataset_name="bpic2012"):  # TODO: 需修改
        self._trace = trace
        self.window_size = window_size

        # 历史预测值序列
        self._prediction_sub_seq = [0] * self.window_size
        self._prediction_sub_seq[-1] = self._trace.prediction_sequence[0]

        # 下一个位置的下标
        self.next_index = 1

        # 构建特征组合
        self.feature = self._build_feature()

        # 采取的action值
        self.actions = [0]

        self._dataset_name = dataset_name
        self._act_error_path = DRLParameters(dataset_name=self._dataset_name).ACTIVITY_ERRORS_PATH

    def _build_feature(self):
        """ 构建特征
        :return:
        """
        feature = []

        # 历史花费时间序列
        cost_time_sub_seq = [0] * self.window_size
        cost_time_sub_seq.extend(self._trace.time_sequence[:self.next_index])
        cost_time_sub_seq = cost_time_sub_seq[-self.window_size:]
        feature.append(cost_time_sub_seq)

        # 历史预测值序列
        feature.append(self._prediction_sub_seq)

        # 历史活动序列
        zero_arr = np.array([[0] * self._trace.activity_space] * self.window_size)
        activity_sub_seq = np.concatenate([zero_arr, self._trace.activity_sequence[:self.next_index, :]], axis=0)
        activity_sub_seq = activity_sub_seq[-self.window_size:, :].T.tolist()
        for arr in activity_sub_seq:
            feature.append(arr)

        # 贷款金额
        # feature.append([self._trace.amount_req] * self.window_size)

        # 上一活动时间差
        last_time_diff_sub_seq = [0] * self.window_size
        last_time_diff_sub_seq.extend(self._trace.last_act_time_diff[:self.next_index])
        last_time_diff_sub_seq = last_time_diff_sub_seq[-self.window_size:]
        feature.append(last_time_diff_sub_seq)

        # 午夜时间差
        midnight_time_diff_sub_seq = [0] * self.window_size
        midnight_time_diff_sub_seq.extend(self._trace.midnight_time_diff[:self.next_index])
        midnight_time_diff_sub_seq = midnight_time_diff_sub_seq[-self.window_size:]
        feature.append(midnight_time_diff_sub_seq)

        # 星期
        zero_arr_7 = np.array([[0] * 7] * self.window_size)
        week_sub_seq = np.concatenate([zero_arr_7, self._trace.weekday[:self.next_index, :]], axis=0)
        week_sub_seq = week_sub_seq[-self.window_size:, :].T.tolist()
        for arr in week_sub_seq:
            feature.append(arr)
        return np.array(feature).T

    def step_to_end(self):
        """ 流程进行到结束
        :return: 下一个状态，奖励值
        """
        while self.next_index < self._trace.trace_length:
            self._prediction_sub_seq.append(self._prediction_sub_seq[-1])
            self.next_index += 1
        self._prediction_sub_seq = self._prediction_sub_seq[-self.window_size:]
        self.feature = self._build_feature()
        return Reward.area_reward(self, 0, alpha=0)  # TODO: 需修改

    def step_next(self, step_length):
        """ 流程进行step_length步长
        :param step_length: 流程进行的长度
        :return: 下一个状态，奖励值
        """
        # reward = Reward.delay_reward(self, step_length, act_error_path=self._act_error_path)  # TODO: 需修改
        reward = Reward.zero_reward(self, step_length)
        self.actions.append(self.next_index-1+step_length)
        for i in range(step_length-1):
            self._prediction_sub_seq.append(self._prediction_sub_seq[-1])
            self.next_index += 1
        self._prediction_sub_seq.append(self._trace.prediction_sequence[self.next_index])
        self._prediction_sub_seq = self._prediction_sub_seq[-self.window_size:]
        self.next_index += 1
        self.feature = self._build_feature()
        return reward

    def clone(self):
        return copy.deepcopy(self)

    def get_trace_ytrue(self):
        return self._trace.true_value

    def get_trace_all_preds(self):
        return self._trace.prediction_sequence

    def get_trace_times(self):
        return self._trace.time_sequence

    def get_trace(self):
        return self._trace

    def get_prediction_sub_sequence(self):
        return self._prediction_sub_seq


if __name__ == '__main__':
    from config import Variables
    pass
